package TAIC.flicker.LM ;

import java.io.File;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Random;
import java.util.Scanner;
import java.util.TreeMap;

public class BayesInDoc {
	
	static Random random = new Random () ; 	
	static double Pic_Min = 0.4 ;  // Threashold for min pictotoriability
	
	private double IG [] ; 
	private double P_w [] ;  // used for IG
	private double NP_w_c [][] ; //  P ( c | not w ), used for IG  
	
	public double P_c [] ;
	public double P_c_w [][] ;  // P(  w | c )
	private double P_w_c [][] ;  // P(  c | w ) 
	private int classes = 2;
	int KeyNumber ;
	TreeMap < String , int [] > map = new TreeMap < String , int [] > () ;   //  int [ classes ]  stores the id( an int ) of a new added word
	public TreeMap < Integer , String > inverseMap = new TreeMap < Integer , String > () ;  // inverseMap stroing the correspondent string of a key(int)
	Stemmer stemmer = new Stemmer ( ) ;
	HashSet < String > validWords = new HashSet < String > () ;

	public BayesInDoc ( String validWordFile, int size ) throws Exception {   
		KeyNumber = size ; 
		Scanner scanner = new Scanner ( new File ( validWordFile )); 
		while ( scanner.hasNext() ) validWords.add( stemmer.stem( scanner.next()));
		scanner.close() ;
	}
		
	public ArrayList < KeyWord > getKey () throws Exception {
		ArrayList < KeyWord > result = null ;
		int i , j ; 
 
		train () ; 		
		caleIG () ;
		
		for ( i = 0 ; i < classes ; i ++ ) 
			if ( i == 0 ) result = getKeyInClass ( i ) ;
			else {
				ArrayList < KeyWord > temp = getKeyInClass ( i ) ;
				for ( j = 0 ; j < temp.size(); j ++ ) result.add ( temp.get ( j ) ) ;
			}	
	
		PrintStream fout = new PrintStream ( new File ( "KeyWords.txt" ) );
		for ( i = 0 ; i < result.size() ; i ++ ) fout.println ( result.get(i).classNo+ " " + result.get ( i ).word + " " + result.get(i).prob );

		return result ; 
	}
	
	public ArrayList < KeyWord > getKeyInClass ( int desiredClass ) {
		if ( KeyNumber > map.size () ) KeyNumber = map.size() - 1 ;
		int i , j;
		
		//System.out.println( IG[ map.get("tire")[classes] ] ) ;
		Node arr [ ] = new Node [ map.size() ] ;
		for ( i = 0 ; i < map.size() ; i ++ ) 
			arr [ i ] = new Node ( i , P_c_w [ desiredClass ][ i ], IG[ i ] ) ;
		
		Arrays.sort ( arr ) ;                   //  sort the arr to get a small query set
		//for (i = 0 ; i < KeyNumber ; i ++ ) System.out.println ( inverseMap.get (arr [ i ].key ));
		
		ArrayList < KeyWord > result = new ArrayList < KeyWord > () ; 		
		
		for ( i = 0 ; i < map.size() ; i ++ ) {
			if ( belongsTo ( arr [ i ].key ) ==desiredClass )
				result.add ( new KeyWord ( inverseMap.get(arr[i].key) , arr[ i ].prob , desiredClass ) ) ;
			if ( result.size () >= KeyNumber ) break ; 
		}		
		
		/* resort using probability */
		for ( i = 0 ;i < result.size() ; i ++ ) {
			for ( j = i + 1 ; j < result.size(); j ++ )
				if ( result.get(i).prob < result.get(j).prob ) {
					KeyWord kw = result.get(i );
					result.set(i , result.get(j));
					result.set(j, kw) ;
				}
		}

		//for ( i = 0 ; i < KeyNumber ; i ++ ) System.out.println( inverseMap.get( arr [ i ].key ) + " " + arr [ i ].IG );
		return result ; 
	}
	
	private void train () throws Exception {
		int docs = 0 , i , j , classNo , totalWords ; 
		String s ; 
		String  word; 
		int tempClass [ ] = null ;
		int classHasWord [] = null ; 
		int arr [ ] ; 
		double sum , nsum; 
		Scanner scanner ;
		
		totalWords = 0 ; 
		Scanner configScanner = new Scanner ( new File ( "temp.txt" ) ) ;
		classes = configScanner.nextInt () ; 
		tempClass = new int [ classes ];
		classHasWord = new int [ classes ] ; 
		P_c = new double [ classes ] ;

		while ( configScanner.hasNextInt() ) {   // Iterate all the document in the training set
			docs ++ ;
			classNo = configScanner.nextInt () ;
			tempClass [ classNo ] ++ ; 

			String filename = configScanner.nextLine().trim() ;
			scanner = new Scanner ( new File (filename) ) ;
			//System.out.println ( filename ) ;

			while ( scanner.hasNext () ) {    //   Iterate all the tagged words in a document 
				word = changeForm( scanner.next () );
				if ( word == null ) continue ;
				word = stemmer.stem( word ) ;
				if ( word.length() == 0 ) continue ;

				if ( validWords.contains( word ) ) {
					addWord ( word , classNo ) ; 
					classHasWord [ classNo ] ++ ;
					totalWords ++ ; 
				}
			}
		}

		int vocabulary = map.size () ; 
		P_w = new double [ vocabulary ] ;
		P_c_w = new double [ classes ][ vocabulary ] ; 
		P_w_c = new double [ vocabulary ][ classes ] ; 
		NP_w_c = new double [ vocabulary ][ classes ] ;
		IG = new double [ vocabulary ] ; 
		for ( i = 0 ; i < classes ; i ++ ) P_c [ i ]= ( double )( tempClass [ i ] ) / docs ; 
		for ( Iterator<String> it = map.keySet ().iterator() ; it.hasNext() ; ) {
			s = it.next () ; 
			arr = map.get ( s ) ;
			if ( arr != null )  {
				for ( i = 0 ; i < classes ; i ++ )  
					P_c_w [ i ][ arr[ classes ] ] = (double)( (map.get ( s ))[i] +  1 )/( classHasWord [ i ] + vocabulary ) ;
			}
		}

		for ( i = 0 ; i < vocabulary ; i ++ ) {
			int [] number = map.get( inverseMap.get(i));
			int total = 0 ; 
			for ( j = 0 ; j < classes ; j ++)  total += number [ j ] ; 
			P_w [ i ] = (double)( total + 1 ) / ( vocabulary + totalWords ) ;
		}

		//PrintStream fout = new PrintStream ( new File ( "words.txt" ) );
		for ( i = 0 ; i < vocabulary ; i ++ ) {
			sum = 1e-9 ; 
			nsum = 1e-9 ; 
			for ( j = 0 ; j < classes ; j ++ ) {
				P_w_c [ i ][ j ] = P_c_w [ j ][ i ] * P_c [ j ] / P_w [ i ];
				NP_w_c [ i ][ j ] = ( 1 - P_c_w [ j ][ i ] ) * P_c [ j ] / (1 - P_w[ i ] ); 
				sum += P_w_c [ i ][ j ] ;
				nsum += NP_w_c[ i ][ j ] ;
			}
			for ( j = 0 ; j < classes ; j ++ ) {
				P_w_c [ i ][ j ] /= sum ;
				NP_w_c [ i ][ j ] /= nsum ; 
			}
			//	fout.println ( P_w_c [ i ][ 0 ] + "|" + P_w_c [ i ][ 1 ] + "|" + inverseMap.get ( i ) ) ; 
		}
		//fout.close();
	}
	

	void caleIG () {
		for ( int i = 0  ; i < map.size() ; i ++ ) {
			double t1 = 0 ;
			double t2 = 0 ; 
			for ( int j = 0 ; j < classes ; j ++ ) {
				t1 += P_w_c [i][j] * Math.log( P_w_c [i][j]) ;
				t2 += NP_w_c [i][j] * Math.log( NP_w_c [i][j]) ;
			}
			t1 *= P_w [ i ];
			t2 *= ( 1 - P_w [ i ]) ;
			IG [ i ] = t1 + t2 ; 
		}
	}
	
	int belongsTo ( int word ){
		int result = 0 ;
		for ( int i = 1 ; i < classes ;i ++  ) 
			if ( P_w_c [ word ][ i ] > P_w_c [ word ][ result ]) result = i; 
		return result ;
	}
	
	private void addWord ( String s , int classNo ) {
		int [] arr = map.get ( s ) ;
		if ( arr == null ) {
			arr = new int [ classes + 1 ] ; 
			arr [ classNo ] = 1 ; 
			arr [ classes ] = map.size () ; 
			inverseMap.put ( arr[ classes ] , s ) ;
			map.put ( s , arr ) ;
		}
		else {
			arr [ classNo ] ++ ; 
		}
	}

	String changeForm ( String s ) {
		if ( s.indexOf( "://" ) != -1 || s.indexOf( "www." ) != -1 || s.indexOf( "<" ) != -1 || s.indexOf( ">" ) != -1) return null ;		
		
		int end = s.length () , start = 0 ;
		while ( start < end && !Character.isLetterOrDigit ( s.charAt ( start ) ) ) start ++ ;
		end -- ; 
		while ( end > start && !Character.isLetterOrDigit ( s.charAt ( end ) ) ) end -- ;
		if ( end > start ) return s.substring ( start , end + 1 ).toLowerCase () ;
		else return null ; 
	}
}

class Node implements Comparable<Node> {
	int key  ;
	double prob ; 
	double IG = 0 ; 
	
	public Node ( int a , double b ) {
		key = a ;
		prob = b ; 
	}
	
	public Node ( int a , double b , double IG ){
		key =a ;
		prob = b ;
		this.IG = IG ; 
	}
	
	public int compareTo ( Node o ) {
		if ( IG > o.IG ) return -1 ;
		else if ( IG == o.IG ) return 0 ;
		else return 1; 
	}
	
}